import os
import re
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import sklearn
from sklearn.model_selection import train_test_split
from sklearn import metrics
import skimage
import skimage.color
import skimage.io
import skimage.feature
import skimage.transform
from sklearn.base import BaseEstimator, TransformerMixin
from skimage.color import rgb2gray
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import GridSearchCV
import sklearn.metrics
from sklearn.preprocessing import StandardScaler

# load the data
data = pickle.load(open('./pickle_files/data_animals_head_20.pickle','rb'))
print("Load pickle data")

dataDesc = data['description']  # There are 20 classes and 2057 images are there. All the images are 80 x 80 (rgb)
X = data['data']                # Contains all images, 
y = data['target']              # Ground truth
labels = data['labels']         # Label names of all dataset

# split the data into train and test
x_train,x_test,y_train,y_test = train_test_split(X,y,test_size=0.2,stratify=y)
print("x_train.shape: {}; x_test.shape: {}; len(y_train): {}; len(y_test): {}".format(x_train.shape,x_test.shape,len(y_train),len(y_test)))

# Feature extraction

img = skimage.io.imread('Images/lion.bmp')
plt.figure(figsize=(6,10))
plt.imshow(img)
plt.show()

# hog transoform for feature extraction.
img_re = skimage.transform.rescale(img,1/3).reshape((302,215))
gray = skimage.color.rgb2gray(img_re)
feature_vector, hog_img = skimage.feature.hog(gray,orientations=9,
    pixels_per_cell=(14,14),
    cells_per_block=(2,2),
    visualize=True)

print("len(feature_vector): ", len(feature_vector))

plt.figure(figsize=(12,10))
plt.subplot(1,2,1)
plt.imshow(gray,cmap='gray')
plt.subplot(1,2,2)
plt.imshow(hog_img,cmap='gray')
plt.show()

# Transform all image to grey.
class rgb2gray_transform(BaseEstimator,TransformerMixin):
    import skimage.color
    def __init__(self):
        pass
    
    def fit(self,X,y=None):
        return self
    
    def transform(self,X,y=None):
        return np.array([skimage.color.rgb2gray(x) for x in X])

rgb2grayobj = rgb2gray_transform()
x_train_gray = rgb2grayobj.fit_transform(x_train)

print('x_train_gray.shape: ', x_train_gray.shape)

# Show RGB image with grey one.
plt.subplot(1,2,1)
plt.imshow(x_train[12])
plt.subplot(1,2,2)
plt.imshow(x_train_gray[12],cmap='gray')
plt.show()

# Extract the feature
class hogtransformer(BaseEstimator,TransformerMixin):
    import skimage.feature
    def __init__(self,orientations=9,pixels_per_cell=(8, 8),cells_per_block=(3, 3),):
        self.orientations = orientations
        self.pixels_per_cell = pixels_per_cell
        self.cells_per_block = cells_per_block
    def fit(self,X,y=None):
        return self
    def transform(self,X,y=None):
        def local_hog(img):
            hog_features= skimage.feature.hog(img,orientations=self.orientations,
                                pixels_per_cell=self.pixels_per_cell,
                                cells_per_block=self.cells_per_block)         
            return hog_features
        hfeatures = np.array([local_hog(x) for x in X])
        return hfeatures

hogt = hogtransformer()
x_train_hog = hogt.fit_transform(x_train_gray)
print('x_train_hog.shape: ',x_train_hog.shape)

# Image classifier
model_sgd = SGDClassifier(loss='hinge',learning_rate='adaptive',
                          early_stopping=True,eta0=0.1,)

grayify = rgb2gray_transform()
hogify = hogtransformer()
scalify = StandardScaler()

# pipeline
# step-1: convert into grayscale
x_train_gray = grayify.fit_transform(x_train)
# step-2: extract the features
x_train_hog = hogify.fit_transform(x_train_gray)
# step-3: Normalization
x_train_scale = scalify.fit_transform(x_train_hog)
# step-4: machine learning
model_sgd.fit(x_train_scale,y_train)

x_test_gray = grayify.fit_transform(x_test)
# step-2: extract the features
x_test_hog = hogify.fit_transform(x_test_gray)
# step-3: Normalization
x_test_scale = scalify.transform(x_test_hog)

y_pred_test = model_sgd.predict(x_test_scale)

print("y_pred_test: ", y_pred_test)

# Model Evaluation
cr = sklearn.metrics.classification_report(y_test,y_pred_test,output_dict=True)
print(pd.DataFrame(cr).T)
print("Model evaluation score: ", metrics.cohen_kappa_score(y_test,y_pred_test))
